package joe;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import static joe.UnsafeUtil.getFloat;
import static joe.UnsafeUtil.putFloat;
import static joe.UnsafeUtil.putLong;

public class Mat4 { // row-major; left-hand coordinate
	private static final long identityConst = UnsafeUtil.allocateMemory(64);

	static {
		UnsafeUtil.zeroMemory(identityConst, 64);
		putFloat(identityConst, 1);
		putFloat(identityConst + 5 * 4, 1);
		putFloat(identityConst + 10 * 4, 1);
		putFloat(identityConst + 15 * 4, 1);
	}

	public final long ptr = Engine.current().getMemPool().alloc64(this);

	public @NotNull Mat4 zero() {
		// UnsafeUtil.zeroMemory(ptr, 64);
		var p = ptr;
		putLong(p, 0);
		putLong(p + 8, 0);
		putLong(p + 16, 0);
		putLong(p + 24, 0);
		putLong(p + 32, 0);
		putLong(p + 40, 0);
		putLong(p + 48, 0);
		putLong(p + 56, 0);
		return this;
	}

	public @NotNull Mat4 identity() {
		UnsafeUtil.copyMemory(identityConst, ptr, 64);
		return this;
	}

	/**
	 * @param i 0..15
	 */
	public float get(int i) {
		return getFloat(ptr + i * 4L);
	}

	/**
	 * @param i 0..15
	 */
	public @NotNull Mat4 set(int i, float v) {
		putFloat(ptr + i * 4L, v);
		return this;
	}

	/**
	 * @param i 0..3
	 */
	public @NotNull Mat4 setRow(int i, float x, float y, float z, float w) {
		putFloat(ptr + i * 16L, x, y, z, w);
		return this;
	}

	/**
	 * @param i 0..3
	 */
	public @NotNull Mat4 setRow(int i, @NotNull Vec4 v) {
		UnsafeUtil.copyMemory(v.ptr, ptr + i * 16L, 16);
		return this;
	}

	public @NotNull Mat4 set(@NotNull Mat4 m) {
		UnsafeUtil.copyMemory(m.ptr, ptr, 64);
		return this;
	}

	public @NotNull Mat4 mul(@NotNull Mat4 m) {
		final var p = ptr;
		UnsafeUtil.matMulAvx(p, m.ptr, p);
		return this;
	}

	public @NotNull Mat4 mul(@NotNull Mat4 m0, @NotNull Mat4 m1) {
		UnsafeUtil.matMulAvx(m0.ptr, m1.ptr, ptr);
		return this;
	}

	public @NotNull Mat4 translate(float x, float y, float z) {
		final var p = ptr;
		putFloat(p + 48, getFloat(p + 48) + x);
		putFloat(p + 52, getFloat(p + 52) + y);
		putFloat(p + 56, getFloat(p + 56) + z);
		return this;
	}

	// m * scale      *  rotate(rz)  *  rotate(rx)  *  rotate(ry)  * translate  => this
	//   | sx  0  0 0 |  cos sin 0 0 | 1   0   0  0 | cos 0 -sin 0 |  1  0  0 0 |
	//   |  0 sy  0 0 | -sin cos 0 0 | 0  cos sin 0 |  0  1   0  0 |  0  1  0 0 |
	//   |  0  0 sz 0 |   0   0  1 0 | 0 -sin cos 0 | sin 0  cos 0 |  0  0  1 0 |
	//   |  0  0  0 1 |   0   0  0 1 | 0   0   0  1 |  0  0   0  1 | tx ty tz 1 |
	public @NotNull Mat4 transform(@Nullable Mat4 m, // null for identity
								   float tx, float ty, float tz, // translate
								   float rx, float ry, float rz, // rotate(degree, counter-clockwise towards the axis)
								   float sx, float sy, float sz) { // scale
		var engine = Engine.current();
		if (m != null) {
			set(m);
			if (sx != 1 || sy != 1 || sz != 1)
				mul(engine.tmpS.set(0, sx).set(5, sy).set(10, sz));
		} else
			zero().set(0, sx).set(5, sy).set(10, sz).set(15, 1);
		if (rz != 0) { // roll
			final var r = rz * (Math.PI / 180);
			final var c = (float)Math.cos(r);
			final var s = (float)Math.sin(r);
			mul(engine.tmpRz.set(0, c).set(1, s).set(4, -s).set(5, c));
		}
		if (rx != 0) { // pitch
			final var r = rx * (Math.PI / 180);
			final var c = (float)Math.cos(r);
			final var s = (float)Math.sin(r);
			mul(engine.tmpRx.set(5, c).set(6, s).set(9, -s).set(10, c));
		}
		if (ry != 0) { // yaw
			final var r = ry * (Math.PI / 180);
			final var c = (float)Math.cos(r);
			final var s = (float)Math.sin(r);
			mul(engine.tmpRy.set(0, c).set(2, -s).set(8, s).set(10, c));
		}
		return translate(tx, ty, tz);
	}

	public @NotNull Mat4 ortho(float width, float height, float near, float far) {
		final var invDepth = -1.0 / (far - near);
		final var p = ptr;
		putFloat(p, 2 / width, 0, 0, 0);
		putFloat(p + 16, 0, 2 / height, 0, 0);
		putFloat(p + 32, 0, 0, (float)(invDepth * -2), 0);
		putFloat(p + 48, 0, 0, (float)((far + near) * invDepth), 1);
		return this;
	}

	/**
	 * @param fov   degree, tan(fov/2) = height/2 / near
	 * @param ratio height / width, e.g. 9/16
	 */
	public @NotNull Mat4 frustum(float fov, float ratio, float near, float far) {
		final var f = 1.0 / Math.tan(fov * (Math.PI / 360));
		final var invDepth = 1.0 / (far - near);
		final var p = ptr;
		putFloat(p, (float)(f * ratio), 0, 0, 0);
		putFloat(p + 16, 0, (float)f, 0, 0);
		putFloat(p + 32, 0, 0, (float)((far + near) * invDepth), 1);
		putFloat(p + 48, 0, 0, (float)(far * near * -2 * invDepth), 0);
		return this;
	}

	public @NotNull Mat4 lookAt(float ex, float ey, float ez, // eye-pos
								float dx, float dy, float dz, // look-dir
								float ux, float uy, float uz) { // up-dir (must be right angle with look-dir)
		double ddx = dx;
		double ddy = dy;
		double ddz = dz;
		var d2 = ddx * ddx + ddy * ddy + ddz * ddz;
		if (Math.abs(d2 - 1) > 1e-6) {
			var f = 1 / Math.sqrt(d2);
			ddx *= f; // dd = normalize(dd)
			ddy *= f;
			ddz *= f;
		}
		double dux = ux;
		double duy = uy;
		double duz = uz;
		d2 = dux * dux + duy * duy + duz * duz;
		if (Math.abs(d2 - 1) > 1e-6) {
			var f = 1 / Math.sqrt(d2);
			dux *= f; // du = normalize(du)
			duy *= f;
			duz *= f;
		}
		final var dsx = duy * ddz - duz * ddy; // ds = cross(du, dd)
		final var dsy = duz * ddx - dux * ddz;
		final var dsz = dux * ddy - duy * ddx;
		final var p = ptr;
		putFloat(p, (float)dsx, (float)dux, (float)ddx, 0);
		putFloat(p + 16, (float)dsy, (float)duy, (float)ddy, 0);
		putFloat(p + 32, (float)dsz, (float)duz, (float)ddz, 0);
		putFloat(p + 48,
				(float)-(dsx * ex + dsy * ey + dsz * ez),
				(float)-(dux * ex + duy * ey + duz * ez),
				(float)-(ddx * ex + ddy * ey + ddz * ez), 1);
		return this;
	}

	@Override
	public String toString() {
		final var p = ptr;
		return String.format("""
						%10.3f %10.3f %10.3f %10.3f
						%10.3f %10.3f %10.3f %10.3f
						%10.3f %10.3f %10.3f %10.3f
						%10.3f %10.3f %10.3f %10.3f
						""",
				getFloat(p), getFloat(p + 4), getFloat(p + 8), getFloat(p + 12),
				getFloat(p + 16), getFloat(p + 20), getFloat(p + 24), getFloat(p + 28),
				getFloat(p + 32), getFloat(p + 36), getFloat(p + 40), getFloat(p + 44),
				getFloat(p + 48), getFloat(p + 52), getFloat(p + 56), getFloat(p + 60));
	}
}
